import os
import cv2
import matplotlib.pyplot as plt
from ultralytics import YOLO
import numpy as np

model = YOLO("/www/wwwroot/AIStomatologynew/best.pt")
print("模型初始化成功")


def predict(img, model=model):
    '''
    :return: list
    '''
    result=None
    #如果model不存在则初始化
    if model==None:
        model=YOLO("/www/wwwroot/AIStomatologynew/best.pt")
        print("模型初始化成功")
    img=np.array(img)

    print("原始图像形状：",img.shape)
    if img.ndim != 3:
        if img.ndim ==2:
            w=int(pow(img.shape[0],0.5))
            h=w
            img=img.reshape((h,w,img.shape[1]))
            print("新图像维度为：",img.shape)
            plt.imshow(img)
            plt.show()
        else:
            print("图像维度错误！")
    r,g,b=cv2.split(img)
    img=cv2.merge([b,g,r])
    # 模型预测，save=True 的时候表示直接保存yolov8的预测结果
    metrics = model.predict(
            img.astype("uint8"),
            save=False,
            save_txt=False,
            save_conf=False,
            project='runs',
            nms=True  # 将图像文件名列表传递给 predict 方法
    )

    for i, m in enumerate(metrics):
        # 获取每个boxes的结果
        boxes = m.boxes
        for box in boxes:
            data=box.data
            data=data.tolist()
            print(data)
            if result is None:
                result=data
            else:
                result.extend(data)

    return result